! ==============================================================================
! This file is part of CUPP (Community Unified Physics Package) since 2023.
! ==============================================================================
! Description:
!
!   This module contains the implementation of the GF (Grell-Freitas) cumulus
!   convective parameterization. It currently does not use the subsidence
!   spreading as in G3.
!
! References:
!
!   - Grell, G. A., and S. E. Freitas, 2014: A scale and aerosol aware stochastic
!     convective parameterization for weather and air quality modeling. Atmos.
!     Chem. Phys., 14, 5233-5250, doi:10.5194/acp-14-5233-2014.
! ==============================================================================

module cu_gf_mod

  use cupp_mesh_mod
  use cu_const_mod
  use atm_types_mod
  use atm_formula_mod

  implicit none

  private

  public cu_gf_init
  public cu_gf_final
  public cu_gf_run

  real(r8), parameter :: eps      = tiny(1.0_r8)
  real(r8), parameter :: max_zup  = 3000.0_r8     ! Maximum height above ground for updraft air can originate (m)
  real(r8), parameter :: max_cap0 = 125.0_r8      ! Maximum capping depth (hPa), larger cap = no convection (inversion)
  real(r8), parameter :: tcwi     = 258.0_r8      ! Critical temperature for water/ice convection (K)

  integer  :: ishallow       = 1
  integer  :: imiddle        = 1
  integer  :: ideep          = 1
  ! 0: average, 1: old g, 2: zws, 3: dhdt
  integer  :: closure_scheme = 1
  real(r8) :: tune_flux      = 1.5_r8

  type gf_data_type
    integer :: id = 0
    type(cupp_mesh_type), pointer :: mesh => null()
    real(r8), allocatable :: entr_rate        (:  ) ! Entrainment rate (s-1)
    real(r8), allocatable :: detr_rate        (:,:) ! Detrainment rate (s-1)
    ! Surface variables
    real(r8), allocatable :: w0               (:  ) ! Convective-scale vertical velocity (m s-1)
    real(r8), allocatable :: dt_ex            (:  ) ! Excess temperature perturbation to increase diurnal forcing
    real(r8), allocatable :: dq_ex            (:  ) ! Excess moisture perturbation to increase diurnal forcing
    real(r8), allocatable :: max_cap          (:  ) ! Maximum capping depth (hPa), larger cap = no convection (inversion
    ! Environmental variables
    real(r8), allocatable :: env_t_lev        (:,:) ! Environmental temperature (K)
    real(r8), allocatable :: env_t_lev_new    (:,:) ! Environmental temperature (K)
    real(r8), allocatable :: env_q_lev        (:,:) ! Environmental specific humidity (kg kg-1)
    real(r8), allocatable :: env_q_lev_new    (:,:) ! Environmental specific humidity (kg kg-1)
    real(r8), allocatable :: env_qes          (:,:) ! Environmental saturated specific humidity (kg kg-1)
    real(r8), allocatable :: env_qes_lev      (:,:) ! Environmental saturated specific humidity (kg kg-1)
    real(r8), allocatable :: env_mse          (:,:) ! Environmental moist static energy (J kg-1)
    real(r8), allocatable :: env_mse_lev      (:,:) ! Environmental moist static energy (J kg-1)
    real(r8), allocatable :: env_smse         (:,:) ! Environmental saturated moist static energy (J kg-1)
    real(r8), allocatable :: env_smse_lev     (:,:) ! Environmental saturated moist static energy (J kg-1)
    real(r8), allocatable :: env_qes_new      (:,:) ! Environmental saturated specific humidity (kg kg-1)
    real(r8), allocatable :: env_qes_lev_new  (:,:) ! Environmental saturated specific humidity (kg kg-1)
    real(r8), allocatable :: env_mse_new      (:,:) ! Environmental moist static energy (J kg-1)
    real(r8), allocatable :: env_mse_lev_new  (:,:) ! Environmental moist static energy (J kg-1)
    real(r8), allocatable :: env_smse_new     (:,:) ! Environmental saturated moist static energy (J kg-1)
    real(r8), allocatable :: env_smse_lev_new (:,:) ! Environmental saturated moist static energy (J kg-1)
    real(r8), allocatable :: env_gamma_lev    (:,:)
    real(r8), allocatable :: env_gamma_lev_new(:,:)
    ! Incloud variables
    real(r8), allocatable :: cld_qw           (:,:) ! Incloud wet mixing ratio of cloud water/ice for radiation (kg kg-1)
    ! ...
    integer , allocatable :: k_max_mse        (:  ) ! Level with largest moist static energy content
    ! Stochastic parameters
    real(r8), allocatable :: rn_mom           (:  ) ! Random number for momentum
    real(r8), allocatable :: rn_vmass         (:  ) ! Random number for vertical mass
    real(r8), allocatable :: rn_clos          (:  ) ! Random number for closure
    ! Tendencies
    real(r8), allocatable :: dptdt            (:,:) ! Potential temperature tendency (K s-1)
    real(r8), allocatable :: dqvdt            (:,:) ! Specific humidity tendency (kg kg-1 s-1)
    real(r8), allocatable :: dudt             (:,:) ! U-wind tendency (m s-2)
    real(r8), allocatable :: dvdt             (:,:) ! V-wind tendency (m s-2)
  contains
    procedure :: init  => gf_data_init
    procedure :: clear => gf_data_clear
    final gf_data_final
  end type gf_data_type

  type(gf_data_type), allocatable :: gf_data(:)

  namelist /cu_gf_params/ &
    ishallow            , &
    imiddle             , &
    ideep               , &
    closure_scheme      , &
    tune_flux

contains

  subroutine cu_gf_init(namelist, mesh)

    character(*), intent(in) :: namelist
    type(cupp_mesh_type), intent(in) :: mesh(:)

    integer iblk

    call cu_gf_final()

    allocate(gf_data(size(mesh)))
    do iblk = 1, size(mesh)
      call gf_data(iblk)%init(mesh(iblk))
    end do

  end subroutine cu_gf_init

  subroutine cu_gf_final()

    if (allocated(gf_data)) deallocate(gf_data)

  end subroutine cu_gf_final

  subroutine cu_gf_run(state, tend, dt_phys, dt_cu)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_cu

    associate (gf => gf_data(state%id))
    gf%rn_mom   = 0.0_r8
    gf%rn_vmass = 0.0_r8
    gf%rn_clos  = 0.0_r8

    select case (ishallow)
    case (1)
      call cu_gf_shallow(state, tend, dt_phys, dt_cu)
    end select

    select case (imiddle)
    case (1)

    end select

    select case (ideep)
    case (1)

    end select
    end associate

  end subroutine cu_gf_run

  subroutine cu_gf_shallow(state, tend, dt_phys, dt_cu)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_cu

    integer icol, ilev
    real(r8) buo_flx

    associate (mesh => state%mesh, gf => gf_data(state%id))
    do icol = 1, mesh%ncol
      gf%entr_rate(icol) = 9.0e-5_r8
      gf%max_cap = max_cap0
    end do
    ! Calculate initial detrainment rate.
    do ilev = 1, mesh%nlev
      do icol = 1, mesh%ncol
        gf%detr_rate(icol,ilev) = gf%entr_rate(icol)
      end do
    end do
    ! ...
    do icol = 1, mesh%ncol
      buo_flx = (state%hfx(icol) / cpd + (rv_o_rd - 1) * state%t_bot(icol) * state%qfx(icol) / lv) / state%rho_bot(icol)
      gf%w0(icol) = max(0.0_r8, tune_flux * 0.41_r8 * buo_flx * state%z(icol,mesh%nlev-1) * g / state%t_bot(icol))
      if (gf%w0(icol) > eps) then
        gf%w0(icol) = 1.2_r8 * gf%w0(icol)**0.3333_r8
        gf%dt_ex(icol) = max(0.0_r8, tune_flux * state%hfx(icol) / (state%rho_bot(icol) * cpd * gf%w0(icol)))
        gf%dq_ex(icol) = max(0.0_r8, tune_flux * state%qfx(icol) / (state%rho_bot(icol) * lv  * gf%w0(icol)))
      end if
      ! For shallow convection closure (Grant 2001)
      gf%w0(icol) = max(0.0_r8, tune_flux * 0.41_r8 * buo_flx * state%z(icol,state%pblk(icol)) * g / state%t(icol,state%pblk(icol)))
      gf%w0(icol) = 1.2_r8 * gf%w0(icol)**0.3333_r8
      ! Check if w0 * rho is correct.
      gf%w0(icol) = gf%w0(icol) * state%rho(icol,state%pblk(icol))
    end do

    ! Calculate environmental moist static energy.
    call cu_gf_env(mesh, state%t, state%p, state%qv, state%z, &
                   gf%env_qes, gf%env_mse, gf%env_smse)
    call cu_gf_env(mesh, state%t_new, state%p, state%qv_new, state%z, &
                   gf%env_qes_new, gf%env_mse_new, gf%env_smse_new)
    call cu_gf_cld(mesh, state%t, state%p, state%qv, state%z, state%zs, &
                   gf%env_qes, gf%env_mse, gf%env_smse, &
                   gf%env_t_lev, gf%env_qes_lev, gf%env_q_lev, &
                   gf%env_mse_lev, gf%env_smse_lev, gf%env_gamma_lev)
    call cu_gf_cld(mesh, state%t_new, state%p, state%qv_new, state%z, state%zs, &
                   gf%env_qes_new, gf%env_mse_new, gf%env_smse_new, &
                   gf%env_t_lev_new, gf%env_qes_lev_new, gf%env_q_lev_new, &
                   gf%env_mse_lev_new, gf%env_smse_lev_new, gf%env_gamma_lev_new)

    ! Determine level with largest moist static energy content.
    do icol = 1, mesh%ncol
      if (state%pblk(icol) > 3) gf%max_cap(icol) = state%p(icol,state%pblk(icol))
      gf%k_max_mse(icol) = maxloc(gf%env_mse(icol,2:), 1)
      gf%k_max_mse(icol) = max(2, gf%k_max_mse(icol))
    end do
    end associate

  end subroutine cu_gf_shallow

  subroutine cu_gf_deep(state, tend, dt_phys, dt_cu)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_cu

  end subroutine cu_gf_deep

  subroutine cu_gf_env(mesh, t, p, q, z, qes, mse, smse)

    type(cupp_mesh_type), intent(in) :: mesh
    real(r8), intent(in   ) :: t   (mesh%ncol_with_halo,mesh%nlev)
    real(r8), intent(in   ) :: p   (mesh%ncol_with_halo,mesh%nlev)
    real(r8), intent(in   ) :: q   (mesh%ncol_with_halo,mesh%nlev)
    real(r8), intent(in   ) :: z   (mesh%ncol_with_halo,mesh%nlev)
    real(r8), intent(  out) :: qes (mesh%ncol_with_halo,mesh%nlev)
    real(r8), intent(  out) :: mse (mesh%ncol_with_halo,mesh%nlev)
    real(r8), intent(  out) :: smse(mesh%ncol_with_halo,mesh%nlev)

    integer icol, ilev
    real(r8) es, dse

    do ilev = 1, mesh%nlev
      do icol = 1, mesh%ncol_with_halo
        es = saturation_vapor_pressure(t(icol,ilev))
        qes(icol,ilev) = saturation_vapor_wet_mixing_ratio(p(icol,ilev), es)
        qes(icol,ilev) = max(q(icol,ilev), max(1.0e-16_r8, qes(icol,ilev)))
        dse = cpd * t(icol,ilev) + g * z(icol,ilev)
        mse(icol,ilev) = dse + lv * q(icol,ilev)
        smse(icol,ilev) = dse + lv * qes(icol,ilev)
        if (mse(icol,ilev) >= smse(icol,ilev)) mse(icol,ilev) = smse(icol,ilev)
      end do
    end do

  end subroutine cu_gf_env

  subroutine cu_gf_cld(mesh, t, p, q, z, zs, qes, mse, smse, t_lev, qes_lev, q_lev, mse_lev, smse_lev, gamma_lev)

    type(cupp_mesh_type), intent(in) :: mesh
    real(r8), intent(in   ) :: t        (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(in   ) :: p        (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(in   ) :: q        (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(in   ) :: z        (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(in   ) :: zs       (mesh%ncol_with_halo            )
    real(r8), intent(in   ) :: qes      (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(in   ) :: mse      (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(in   ) :: smse     (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(  out) :: t_lev    (mesh%ncol_with_halo,mesh%nlev  )
    real(r8), intent(  out) :: qes_lev  (mesh%ncol_with_halo,mesh%nlev+1)
    real(r8), intent(  out) :: q_lev    (mesh%ncol_with_halo,mesh%nlev+1)
    real(r8), intent(  out) :: mse_lev  (mesh%ncol_with_halo,mesh%nlev+1)
    real(r8), intent(  out) :: smse_lev (mesh%ncol_with_halo,mesh%nlev+1)
    real(r8), intent(  out) :: gamma_lev(mesh%ncol_with_halo,mesh%nlev+1)

    integer icol, ilev

    do ilev = 2, mesh%nlev
      do icol = 1, mesh%ncol_with_halo
        t_lev    (icol,ilev) = 0.5_r8 * (t   (icol,ilev-1) + t   (icol,ilev))
        qes_lev  (icol,ilev) = 0.5_r8 * (qes (icol,ilev-1) + qes (icol,ilev))
        q_lev    (icol,ilev) = 0.5_r8 * (q   (icol,ilev-1) + q   (icol,ilev))
        mse_lev  (icol,ilev) = 0.5_r8 * (mse (icol,ilev-1) + mse (icol,ilev))
        smse_lev (icol,ilev) = 0.5_r8 * (smse(icol,ilev-1) + smse(icol,ilev))
        mse_lev  (icol,ilev) = min(mse_lev(icol,ilev), smse_lev(icol,ilev))
        gamma_lev(icol,ilev) = lv / cpd * (lv / (rv * t_lev(icol,ilev)**2)) * qes_lev(icol,ilev)
      end do
    end do
    ilev = mesh%nlev + 1
    do icol = 1, mesh%ncol_with_halo
      t_lev    (icol,ilev) = t  (icol,ilev-1)
      qes_lev  (icol,ilev) = qes(icol,ilev-1)
      q_lev    (icol,ilev) = q  (icol,ilev-1)
      mse_lev  (icol,ilev) = cpd * t(icol,ilev-1) + g * zs(icol) + lv * q  (icol,ilev-1)
      smse_lev (icol,ilev) = cpd * t(icol,ilev-1) + g * zs(icol) + lv * qes(icol,ilev-1)
      gamma_lev(icol,ilev) = lv / cpd * (lv / (rv * t_lev(icol,ilev)**2)) * qes_lev(icol,ilev)
    end do

  end subroutine cu_gf_cld

  subroutine gf_data_init(this, mesh)

    class(gf_data_type), intent(inout) :: this
    type(cupp_mesh_type), intent(in), target :: mesh

    call this%clear()

    this%id = mesh%id
    this%mesh => mesh

    allocate(this%entr_rate       (mesh%ncol_with_halo))
    allocate(this%detr_rate       (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%w0              (mesh%ncol_with_halo))
    allocate(this%dt_ex           (mesh%ncol_with_halo))
    allocate(this%dq_ex           (mesh%ncol_with_halo))
    allocate(this%max_cap         (mesh%ncol_with_halo))
    allocate(this%env_qes         (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%env_qes_lev     (mesh%ncol_with_halo,mesh%nlev+1))
    allocate(this%env_mse         (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%env_mse_lev     (mesh%ncol_with_halo,mesh%nlev+1))
    allocate(this%env_smse        (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%env_smse_lev    (mesh%ncol_with_halo,mesh%nlev+1))
    allocate(this%env_qes_new     (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%env_qes_lev_new (mesh%ncol_with_halo,mesh%nlev+1))
    allocate(this%env_mse_new     (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%env_mse_lev_new (mesh%ncol_with_halo,mesh%nlev+1))
    allocate(this%env_smse_new    (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%env_smse_lev_new(mesh%ncol_with_halo,mesh%nlev+1))
    allocate(this%cld_qw          (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%k_max_mse       (mesh%ncol_with_halo))
    allocate(this%rn_mom          (mesh%ncol_with_halo))
    allocate(this%rn_vmass        (mesh%ncol_with_halo))
    allocate(this%rn_clos         (mesh%ncol_with_halo))

  end subroutine gf_data_init

  subroutine gf_data_clear(this)

    class(gf_data_type), intent(inout) :: this

    this%id   = 0
    this%mesh => null()

    if (allocated(this%entr_rate       )) deallocate(this%entr_rate       )
    if (allocated(this%detr_rate       )) deallocate(this%detr_rate       )
    if (allocated(this%w0              )) deallocate(this%w0              )
    if (allocated(this%dt_ex           )) deallocate(this%dt_ex           )
    if (allocated(this%dq_ex           )) deallocate(this%dq_ex           )
    if (allocated(this%max_cap         )) deallocate(this%max_cap         )
    if (allocated(this%env_qes         )) deallocate(this%env_qes         )
    if (allocated(this%env_qes_lev     )) deallocate(this%env_qes_lev     )
    if (allocated(this%env_mse         )) deallocate(this%env_mse         )
    if (allocated(this%env_mse_lev     )) deallocate(this%env_mse_lev     )
    if (allocated(this%env_smse        )) deallocate(this%env_smse        )
    if (allocated(this%env_smse_lev    )) deallocate(this%env_smse_lev    )
    if (allocated(this%env_qes_new     )) deallocate(this%env_qes_new     )
    if (allocated(this%env_qes_lev_new )) deallocate(this%env_qes_lev_new )
    if (allocated(this%env_mse_new     )) deallocate(this%env_mse_new     )
    if (allocated(this%env_mse_lev_new )) deallocate(this%env_mse_lev_new )
    if (allocated(this%env_smse_new    )) deallocate(this%env_smse_new    )
    if (allocated(this%env_smse_lev_new)) deallocate(this%env_smse_lev_new)
    if (allocated(this%cld_qw          )) deallocate(this%cld_qw          )
    if (allocated(this%k_max_mse       )) deallocate(this%k_max_mse       )
    if (allocated(this%rn_mom          )) deallocate(this%rn_mom          )
    if (allocated(this%rn_vmass        )) deallocate(this%rn_vmass        )
    if (allocated(this%rn_clos         )) deallocate(this%rn_clos         )

  end subroutine gf_data_clear

  subroutine gf_data_final(this)

    type(gf_data_type), intent(inout) :: this

    call this%clear()

  end subroutine gf_data_final

end module cu_gf_mod